import copy
from typing import Union

import dspy
import re
from .storm_dataclass import StormArticle
from ...interface import ArticlePolishingModule
from ...utils import ArticleTextProcessing
from .savetemp import savetemp


class StormArticlePolishingModule(ArticlePolishingModule):
    """
    The interface for article generation stage. Given topic, collected information from
    knowledge curation stage, generated outline from outline generation stage.
    """

    def __init__(
        self,
        article_gen_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel],
        article_polish_lm: Union[dspy.dsp.LM, dspy.dsp.HFModel],
    ):
        self.article_gen_lm = article_gen_lm
        self.article_polish_lm = article_polish_lm

        self.polish_page = PolishPageModule(
            write_lead_engine=self.article_gen_lm, polish_engine=self.article_polish_lm
        )

    def polish_article(
        self, topic: str, draft_article: StormArticle, remove_duplicate: bool = False
    ) -> StormArticle:
        """
        Polish article.

        Args:
            topic (str): The topic of the article.
            draft_article (StormArticle): The draft article.
            remove_duplicate (bool): Whether to use one additional LM call to remove duplicates from the article.
        """

        article_text = draft_article.to_string()
        polish_result = self.polish_page(
            topic=topic, draft_page=article_text, polish_whole_page=remove_duplicate
        )
        
        polished_article = polish_result.page
        polished_article_dict = ArticleTextProcessing.parse_article_into_dict(
            polished_article
        )

        
        polished_article = copy.deepcopy(draft_article)
        polished_article.insert_or_create_section(article_dict=polished_article_dict)


        #进行引用格式替换
        def replace_citations(content: str) -> str:
            """
            Replace [1], [2], ... with <sup>1</sup>, <sup>2</sup>, ...
            This operation strictly applies only to the content string.
            """
            if content:
                return re.sub(r"\[(\d+)\]", r"<sup>\1</sup>", content)
            return content

        def replace_references_only(node):
            """
            Replace citation formats in the content of each section without altering the structure.

            Args:
                node: The root node of the article's section tree.
            """
            if not node:
                return
            
            # Replace citations in the current node's content
            if node.content:
                node.content = replace_citations(node.content)
            
            # Recursively apply to all children nodes
            for child in node.children:
                replace_references_only(child)

        # 调用替换函数，仅对内容执行替换
        replace_references_only(polished_article.root)


        # 生成 # References 部分
        references_section = self.generate_references_section(draft_article)
        if references_section:
            polished_article.update_section(
                current_section_content=references_section,
                current_section_info_list=None,  # 不更新引用，只添加内容
                parent_section_name=None,  # 添加到根节点
            )
        savetemp("article_polish/raw_lead_section.txt", polish_result.lead_section)
        #加入LeadSection
        # import re
        # if re.match(r'(?i)^abstract?',polish_result.lead_section.split("\n")[0].strip("#").strip()):
        #     lead_section = polish_result.lead_section
        abstract_title = polish_result.lead_section.split("\n")[0].strip("#").strip().split(" ")
        if len(abstract_title) == 1:
            first_line_removed = polish_result.lead_section.split('\n', 1)[1]
            lead_section = f"**Abstract**\n{first_line_removed}"
        else:
            lead_section = f"**Abstract**\n{polish_result.lead_section.strip('#').strip()}"

        savetemp("article_polish/lead_section.txt", lead_section)

        polished_article = polished_article.to_string()

        combined_article = re.sub(r"(#.*?\n)(.*?)(?=\n##|\Z)", rf"\1{lead_section}", polished_article, count=1, flags=re.DOTALL)

        # combined_article = f"{lead_section}\n{polished_article}"  # 确保合并时中间有两个换行分隔
        
        savetemp("article_polish/combined_article.txt", combined_article)

        polished_article_dict = ArticleTextProcessing.parse_article_into_dict(
            combined_article
        )
        
        polished_article = copy.deepcopy(draft_article)
        polished_article.insert_or_create_section(article_dict=polished_article_dict)
        
        #后续处理
        polished_article.post_processing()
        savetemp("article_polish/polished_article.json", polished_article)
        return polished_article

    #为文章末尾增加排序好的参考文献
    def generate_references_section(self, article: StormArticle) -> str:
        """
        Generate a References section from the article's URL-to-info mapping.

        Args:
            article (StormArticle): The article containing reference information.
    
        Returns:
            str: The formatted References section as a string.
        """
        references = article.reference["url_to_info"]
        if not references:
            return ""
    
        # 使用统一索引的顺序生成引用内容
        unified_index_to_url = {
            index: url
            for url, index in article.reference["url_to_unified_index"].items()
        }
    
        # 按索引排序并生成引用内容
        sorted_references = sorted(unified_index_to_url.items())  # 按索引排序
        references_list = []
        for index, url in sorted_references:
            info = references[url]
            title = info.title if info.title else "No Title Available"
            description = info.description

            # 确保标题后面有句号
            if not title.endswith("."):
                title += "."
            
            references_list.append(f"[{index}] {title}{description}")
    
        references_content = "# References\n" + "\n".join(references_list)
        return references_content



class WriteLeadSection(dspy.Signature):
    """
    Write an abstract for a structured long literature review with the following guidelines:
    1. The abstract should provide a concise overview of the topic, including its significance, current state of research, and any key challenges or trends.
    2. The abstract should stand on its own, summarizing the main points from the review without requiring additional context.
    3. The abstract should be concise and contain no more than four well-composed paragraphs.
    4. No Citations,No References in this section!
    5.Section name is #Abstract
    """

    topic = dspy.InputField(prefix="The topic of the page: ", format=str)
    draft_page = dspy.InputField(prefix="The draft page:\n", format=str)
    lead_section = dspy.OutputField(prefix="Write the lead section:\n", format=str)


class PolishPage(dspy.Signature):
    """
    You are a faithful academic writing expert that is good at finding repeated information in the article and deleting them to make sure there is no repetition in the article.
    You won't delete any non-repeated part in the article. You will keep the inline citations and article structure (indicated by "#", "##", etc.) appropriately.
    Do your job for the following article.
    """

    draft_page = dspy.InputField(prefix="The draft article:\n", format=str)
    page = dspy.OutputField(prefix="Your revised article:\n", format=str)


class PolishPageModule(dspy.Module):
    def __init__(
        self,
        write_lead_engine: Union[dspy.dsp.LM, dspy.dsp.HFModel],
        polish_engine: Union[dspy.dsp.LM, dspy.dsp.HFModel],
    ):
        super().__init__()
        self.write_lead_engine = write_lead_engine
        self.polish_engine = polish_engine
        self.write_lead = dspy.Predict(WriteLeadSection)
        self.polish_page = dspy.Predict(PolishPage)

    def forward(self, topic: str, draft_page: str, polish_whole_page: bool = True):
        # NOTE: Change show_guidelines to false to make the generation more robust to different LM families.
        with dspy.settings.context(lm=self.write_lead_engine, show_guidelines=False):
            lead_section = self.write_lead(
                topic=topic, draft_page=draft_page
            ).lead_section
            if "The lead section:" in lead_section:
                lead_section = lead_section.split("The lead section:")[1].strip()
        if polish_whole_page:
            # NOTE: Change show_guidelines to false to make the generation more robust to different LM families.
            with dspy.settings.context(lm=self.polish_engine, show_guidelines=False):
                page = self.polish_page(draft_page=draft_page).page
        else:
            page = draft_page

        return dspy.Prediction(lead_section=lead_section, page=page)
